#include "NumberTheory.h"
#include "Sequence.h"
#include <algorithm>
#include <ctime>
#include <sstream>


Integer NumberTheory::GCD(const Integer &x, const Integer &y)
{
	if (x == 0)
		return std::abs(y);
	else if (y == 0)
		return std::abs(x);

	if (x < 0 || y < 0)
		return GCD(std::abs(x), std::abs(y));

	if (x > y)
		return GCD(x % y, y);
	else
		return GCD(y % x, x);
}

Integer NumberTheory::LCM(const Integer &x, const Integer &y)
{
	if (x == 0 || y == 0)
		return 0;

	return std::abs(x / GCD(x, y) * y);
}

Integer NumberTheory::ExGCD(const Integer &x, const Integer &y, Integer &a, Integer &b)
{
	if (x == 0)
	{
		a = 0;
		b = std::sgn(y);
		return std::abs(y);
	}
	else if (y == 0)
	{
		a = std::sgn(x);
		b = 0;
		return std::abs(x);
	}

	if (std::abs(x) > std::abs(y))
		return ExGCD(y, x, b, a);

	a = 0;
	b = 1;
	Integer d(std::abs(y));
	Integer a1(1), b1(0), d1(std::abs(x));

	while (d1 > 0)
	{
		Integer a2 = std::move(a1);
		Integer b2 = std::move(b1);
		Integer d2 = std::move(d1);
		Integer q = Divide(d, d2, d1);
		a1 = a - q * a2;
		b1 = b - q * b2;
		a = std::move(a2);
		b = std::move(b2);
		d = std::move(d2);
	}

	if (x < 0)
		a *= -1;

	if (y < 0)
		b *= -1;

	return d;
}

Polynomial<Integer> NumberTheory::GCD(const Polynomial<Integer> &poly1, const Polynomial<Integer> &poly2)
{
	Polynomial<Rational> p1, p2;
	p1.assign(poly1.begin(), poly1.end());
	p2.assign(poly2.begin(), poly2.end());

	const Polynomial<Rational> p = GCD(p1, p2);
	const unsigned size = p.size();

	if (size <= 1)
		return Polynomial<Integer>(size);

	Integer gcd, lcm(1);
	Polynomial<Integer> poly;

	for (unsigned i = 0; i < size; ++i)
	{
		gcd = GCD(gcd, p[i].Numerator());
		lcm = LCM(lcm, p[i].Denominator());
	}

	for (unsigned i = 0; i < size; ++i)
		poly.push_back((p[i].Numerator() / gcd) * (lcm / p[i].Denominator()));

	return poly;
}

int NumberTheory::Legendre_Jacobi_Kronecker(const Integer &x, const Integer &y)
{
	if (x == 0)
		return y == 1 || y == -1;
	else if (y == 0)
		return x == 1 || x == -1;
	else if (y == 1)
		return 1;
	else if (y < 0)
		return Legendre_Jacobi_Kronecker(x, -y) * std::sgn(x);
	else if (std::IsEven(y))
	{
		if (std::IsEven(x))
			return 0;

		unsigned d = 0;

		while (!y.Bit(++d));

		const int n = ((!(x.Bit(1) ^ x.Bit(2))) || std::IsEven(d)) ? 1 : -1;
		return Legendre_Jacobi_Kronecker(x, y >> d) * n;
	}
	else if (x < 0)
	{
		const int n = y.Bit(1) ? -1 : 1;
		return Legendre_Jacobi_Kronecker(-x, y) * n;
	}
	else if (std::IsEven(x))
	{
		unsigned d = 0;

		while (!x.Bit(++d));

		const int n = ((!(y.Bit(1) ^ y.Bit(2))) || std::IsEven(d)) ? 1 : -1;
		return Legendre_Jacobi_Kronecker(x >> d, y) * n;
	}
	else if (x >= y)
		return Legendre_Jacobi_Kronecker(x % y, y);
	else
	{
		const int n = (x.Bit(1) && y.Bit(1)) ? -1 : 1;
		return Legendre_Jacobi_Kronecker(y % x, x) * n;
	}
}

bool RabinMillerPrimaltyTest(const Integer &n, const Integer &base)
{
	unsigned s = 0;

	while (!n.Bit(++s));

	Residue a = std::Exp(Residue(base, n), n >> s);

	if (a == 1)
		return true;

	for (unsigned k = 0; k < s; ++k, a *= a)
		if (a == n - 1)
			return true;

	return false;
}

bool NumberTheory::IsProbablyPrime(const Integer &n)
{
	if (n == 0 || n == 1)
		return false;
	else if (n < 0)
		return IsProbablyPrime(-n);

	const unsigned smallprime[25] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47,
									 53, 59, 61, 67, 71, 73, 79, 83, 89, 97
									};

	for (int i = 0; i < 25; ++i)
	{
		if (n == smallprime[i])
			return true;
		else if (n % smallprime[i] == 0)
			return false;
	}

	if (!RabinMillerPrimaltyTest(n, 2) || !RabinMillerPrimaltyTest(n, 3))
		return false;

	static const bool q11[11] = {true, true, false, true, true, true, false, false, false, true, false};
	static const bool q63[63] = {true, true, false, false, true, false, false, true, false, true,
								 false, false, false, false, false, false, true, false, true, false,
								 false, false, true, false, false, true, false, false, true, false,
								 false, false, false, false, false, false, true, true, false, false,
								 false, false, false, true, false, false, true, false, false, true,
								 false, false, false, false, false, false, false, false, true, false,
								 false, false, false
								};
	static const bool q64[64] = {true, true, false, false, true, false, false, false, false, true,
								 false, false, false, false, false, false, true, true, false, false,
								 false, false, false, false, false, true, false, false, false, false,
								 false, false, false, true, false, false, true, false, false, false,
								 false, true, false, false, false, false, false, false, false, true,
								 false, false, false, false, false, false, false, true, false, false,
								 false, false, false, false,
								};
	static const bool q65[65] = {true, true, false, false, true, false, false, false, false, true,
								 true, false, false, false, true, false, true, false, false, false,
								 false, false, false, false, false, true, true, false, false, true,
								 true, false, false, false, false, true, true, false, false, true,
								 true, false, false, false, false, false, false, false, false, true,
								 false, true, false, false, false, true, true, false, false, false,
								 false, true, false, false, true
								};

	if (q64[_unsigned(n) & 0x3f])
	{
		const unsigned n0 = n % 45045;

		if (q63[n0 % 63] && q65[n0 % 65] && q11[n0 % 11])
		{
			const Integer sqrtN = IntegerRoot(n, 2);

			if (sqrtN * sqrtN == n)
				return false;
		}
	}

	Integer d(5);
	int m = 0;

	while (m != -1)
	{
		m = NumberTheory::Legendre_Jacobi_Kronecker(d, n);

		if (m == 0)
			return false;
		else if (m == 1)
			d = -d - 2 * std::sgn(d);
	}

	return Sequence::LucasU(n + 1, Residue(1, n), Residue((d - 1) >> 2, n)) == 0;
}

bool NumberTheory::IsPrime(const Integer &n)
{
	return IsProbablyPrime(n);
}

bool NumberTheory::IsComposite(const Integer &n)
{
	if (n == 0)
		return true;
	else if (n == 1 || n == -1)
		return false;
	else
		return !IsProbablyPrime(n);
}

Integer PollardRhoFindFactor(const Integer &n)
{
	for (unsigned a = 1; ; ++a)
	{
		Integer x;
		Integer y(a + 1);

		for (unsigned k = 1; ; k <<= 1)
		{
			x = y;

			for (unsigned i = 0; i < k; ++i)
				y = (y * y + a) % n;

			for (unsigned i = 0; i < k; ++i)
			{
				y = (y * y + a) % n;
				const Integer gcd = NumberTheory::GCD(x - y, n);

				if (gcd > 1)
				{
					if (gcd != n)
						return gcd;
					else
						goto end;
				}
			}
		}

end:
		;
	}
}

Integer SQUFOFFindFactor(const Integer &n)
{
	const Integer sqrtN = IntegerRoot(n, 2);

	if (sqrtN * sqrtN == n)
		return sqrtN;

	static const bool q11[11] = {true, true, false, true, true, true, false, false, false, true, false};
	static const bool q63[63] = {true, true, false, false, true, false, false, true, false, true,
								 false, false, false, false, false, false, true, false, true, false,
								 false, false, true, false, false, true, false, false, true, false,
								 false, false, false, false, false, false, true, true, false, false,
								 false, false, false, true, false, false, true, false, false, true,
								 false, false, false, false, false, false, false, false, true, false,
								 false, false, false
								};
	static const bool q64[64] = {true, true, false, false, true, false, false, false, false, true,
								 false, false, false, false, false, false, true, true, false, false,
								 false, false, false, false, false, true, false, false, false, false,
								 false, false, false, true, false, false, true, false, false, false,
								 false, true, false, false, false, false, false, false, false, true,
								 false, false, false, false, false, false, false, true, false, false,
								 false, false, false, false,
								};
	static const bool q65[65] = {true, true, false, false, true, false, false, false, false, true,
								 true, false, false, false, true, false, true, false, false, false,
								 false, false, false, false, false, true, true, false, false, true,
								 true, false, false, false, false, true, true, false, false, true,
								 true, false, false, false, false, false, false, false, false, true,
								 false, true, false, false, false, true, true, false, false, false,
								 false, true, false, false, true
								};

	for (Integer N(n); ; N += n)
	{
		Integer m = IntegerRoot(N, 2);
		Integer m1, p0;
		Integer p1(m);
		Integer q0(1);
		Integer q1 = N - m * m;

		do
		{
			std::swap(p0, p1);
			std::swap(q0, q1);
			Integer r;
			const Integer b = Divide(m + p0, q0, r);
			p1 = m - r;
			q1 += b * (p0 - p1);

			if (!q64[_unsigned(q1) & 0x3f])
				continue;

			const unsigned n0 = q1 % 45045;

			if (!q63[n0 % 63] || !q65[n0 % 65] || !q11[n0 % 11])
				continue;

			m1 = IntegerRoot(q1, 2);
		}
		while (m1 * m1 != q1);

		p1 = m - (m - p1) % m1;
		Integer q2 = (N - p1 * p1) / m1;

		do
		{
			std::swap(p0, p1);
			std::swap(q0, q2);
			Integer r;
			const Integer b = Divide(m + p0, q0, r);
			p1 = m - r;
			q2 += b * (p0 - p1);
		}
		while (p1 != p0);

		const Integer gcd = NumberTheory::GCD(n, p0);

		if (gcd != 1)
			return gcd;
	}
}

Integer NumberTheory::FindFactor(const Integer &n)
{
	if (n == 0)
		return 0;
	else if (n < 0)
		return FindFactor(-n);

	if (IsPrime(n))
		return n;

	const unsigned smallprime[25] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47,
									 53, 59, 61, 67, 71, 73, 79, 83, 89, 97
									};

	for (unsigned i = 0; i < 25; i++)
		if (n % smallprime[i] == 0)
			return smallprime[i];

	return PollardRhoFindFactor(n);
}

std::vector<std::pair<Integer, unsigned>> NumberTheory::Factor(const Integer &n)
{
	if (n == 0)
		return std::vector<std::pair<Integer, unsigned>>(1, std::make_pair(0, 1));
	else if (n < 0)
	{
		auto factor = Factor(-n);
		factor.insert(factor.begin(), 1, std::make_pair(-1, 1));
		return factor;
	}
	else if (n == 1)
		return std::vector<std::pair<Integer, unsigned>>(1, std::make_pair(1, 1));
	else if (IsPrime(n))
		return std::vector<std::pair<Integer, unsigned>>(1, std::make_pair(n, 1));

	const Integer p = FindFactor(n);
	unsigned k = 0xffffffff;
	Integer n1, n2(n), remainder;

	do
	{
		std::swap(n1, n2);
		n2 = Divide(n1, p, remainder);
		++k;
	}
	while (remainder == 0);

	auto f = Factor(p);

	std::for_each(f.begin(), f.end(),
				  [](std::pair<Integer, unsigned> &x){x = std::make_pair(x.first, x.second);});

	if (n1 == 1)
		return f;

	const auto f2 = Factor(n1);
	std::vector<std::pair<Integer, unsigned>> factor;
	auto iter1 = f.begin();
	auto iter2 = f2.begin();
	const auto end1 = f.end();
	const auto end2 = f2.end();

	while (iter1 != end1 && iter2 != end2)
	{
		if (iter1->first == iter2->first)
		{
			factor.push_back(std::make_pair(iter1->first, iter1->second + iter2->second));
			++iter1;
			++iter2;
		}
		else if (iter1->first < iter2->first)
			factor.push_back(*iter1++);
		else
			factor.push_back(*iter2++);
	}

	while (iter1 != end1)
		factor.push_back(*iter1++);

	while (iter2 != end2)
		factor.push_back(*iter2++);

	return factor;
}

std::vector<std::pair<Integer, int>> NumberTheory::Factor(const Rational &r)
{
	if (r.Denominator() == 1)
	{
		const auto p = Factor(r.Numerator());
		const unsigned size = p.size();
		std::vector<std::pair<Integer, int>> factor;

		for (unsigned i = 0; i < size; ++i)
			factor.push_back(std::make_pair(p[i].first, static_cast<int>(p[i].second)));

		return factor;
	}
	else if (r.Numerator() == 1)
	{
		const auto p = Factor(r.Denominator());
		const unsigned size = p.size();
		std::vector<std::pair<Integer, int>> factor;

		for (unsigned i = 0; i < size; ++i)
			factor.push_back(std::make_pair(p[i].first, -static_cast<int>(p[i].second)));

		return factor;
	}
	else if (r.Numerator() == -1)
	{
		const auto p = Factor(r.Denominator());
		const unsigned size = p.size();
		std::vector<std::pair<Integer, int>> factor(1, std::make_pair(-1, 1));

		for (unsigned i = 0; i < size; ++i)
			factor.push_back(std::make_pair(p[i].first, -static_cast<int>(p[i].second)));

		return factor;
	}
	else
	{
		const auto p = Factor(r.Numerator());
		const auto q = Factor(r.Denominator());
		std::vector<std::pair<Integer, int>> factor;

		if (r < 0)
			factor.push_back(std::make_pair(-1, 1));

		auto iter1 = p.begin();
		auto iter2 = q.begin();
		const auto end1 = p.end();
		const auto end2 = q.end();

		do
		{
			if (iter1->first == iter2->first)
			{
				factor.push_back(std::make_pair(iter1->first, iter1->second + iter2->second));
				++iter1;
				++iter2;
			}
			else if (iter1->first < iter2->first)
			{
				factor.push_back(std::make_pair(iter1->first, static_cast<int>(iter1->second)));
				++iter1;
			}
			else
			{
				factor.push_back(std::make_pair(iter2->first, -static_cast<int>(iter2->second)));
				++iter2;
			}
		}
		while (iter1 != end1 || iter2 != end1);

		while (iter1 != end1)
		{
			factor.push_back(std::make_pair(iter1->first, static_cast<int>(iter1->second)));
			++iter1;
		}

		while (iter2 != end2)
		{
			factor.push_back(std::make_pair(iter2->first, -static_cast<int>(iter2->second)));
			++iter2;
		}

		return factor;
	}
}

std::vector<Polynomial<ResidueField>> EqualDegreeFactor(const Polynomial<ResidueField> &poly, const unsigned degree)
{
	static std::mt19937 randgen(time(nullptr));
	const Integer mod = poly.LeadingCoefficient().Modulus();
	std::vector<Polynomial<ResidueField>> factor;

	if (Degree(poly) == degree)
	{
		const ResidueField invlc = Inverse(poly.LeadingCoefficient());
		Polynomial<ResidueField> p;

		for (const auto &x: poly)
			p.push_back(x * invlc);

		factor.push_back(p);
	}
	else if (mod == 2)
	{
		while (true)
		{
			static std::uniform_int_distribution<unsigned> uid(0, 1);
			Polynomial<ResidueField> a;

			for (unsigned i = 0; i < poly.size(); ++i)
				a.push_back(ResidueField(uid(randgen), 2, true));

			while (!a.empty() && a.back() == 0)
				a.pop_back();

			if (!Degree(a))
				continue;

			Polynomial<ResidueField> g = NumberTheory::GCD(a, poly);

			if (g != Polynomial<ResidueField>(ResidueField(1, 2)) && g != poly)
			{
				factor = EqualDegreeFactor(g, degree);
				const auto f = EqualDegreeFactor(poly / g, degree);
				factor.insert(factor.end(), f.begin(), f.end());
				break;
			}

			Polynomial<ResidueField> b, d(a);

			for (unsigned i = 0; i < degree; ++i)
			{
				b = (b + d) % poly;
				Polynomial<ResidueField> d0;
				d0.push_back(d[0]);

				for (unsigned i = 1; i < d.size(); ++i)
				{
					d0.push_back(ResidueField(0, 2, true));
					d0.push_back(d[i]);
				}

				d = d0 % poly;
			}

			g = NumberTheory::GCD(b, poly);

			if (g != Polynomial<ResidueField>(ResidueField(1, 2)) && g != poly)
			{
				factor = EqualDegreeFactor(g, degree);
				const auto f = EqualDegreeFactor(poly / g, degree);
				factor.insert(factor.end(), f.begin(), f.end());
				break;
			}
		}
	}
	else
	{
		while (true)
		{
			static std::uniform_int_distribution<unsigned> uid(0, _unsigned(mod) - 1);
			static std::uniform_int_distribution<unsigned> uid0;
			Polynomial<ResidueField> a;

			if (mod < 0xffffffff)
				for (unsigned i = 0; i < poly.size(); ++i)
					a.push_back(ResidueField(uid(randgen), mod, true));
			else
				for (unsigned i = 0; i < poly.size(); ++i)
					a.push_back(ResidueField(uid0(randgen), mod, true));

			while (!a.empty() && a.back() == 0)
				a.pop_back();

			if (!Degree(a))
				continue;

			Polynomial<ResidueField> g = NumberTheory::GCD(a, poly);

			if (g != Polynomial<ResidueField>(ResidueField(1, mod, true)) && g != poly)
			{
				factor = EqualDegreeFactor(g, degree);
				const auto f = EqualDegreeFactor(poly / g, degree);
				factor.insert(factor.end(), f.begin(), f.end());
				break;
			}

			auto b = std::Exp(PolynomialResidue<ResidueField>(a, poly), std::Exp(mod, degree) >> 1).Value();
			g = NumberTheory::GCD(b - Polynomial<ResidueField>(ResidueField(1, mod, true)), poly);

			if (g != Polynomial<ResidueField>(ResidueField(1, mod, true)) && g != poly)
			{
				factor = EqualDegreeFactor(g, degree);
				const auto f = EqualDegreeFactor(poly / g, degree);
				factor.insert(factor.end(), f.begin(), f.end());
				break;
			}
		}
	}

	std::sort(factor.begin(), factor.end());
	return factor;
}

std::vector<std::pair<Polynomial<ResidueField>, unsigned>> NumberTheory::Factor(const Polynomial<ResidueField> &poly)
{
	if (!Degree(poly))
	{
		std::vector<std::pair<Polynomial<ResidueField>, unsigned>> factor;

		if (poly == Polynomial<ResidueField>(0))
			factor.push_back(std::make_pair(Polynomial<ResidueField>(), 1));
		else
			factor.push_back(std::make_pair(Polynomial<ResidueField>(1), 1));

		return factor;
	}

	const Integer p = poly.LeadingCoefficient().Modulus();
	const auto squarefreepoly = poly / GCD(poly, Differential(poly));

	if (!Degree(squarefreepoly))
	{
		std::vector<std::pair<Polynomial<ResidueField>, unsigned>> factor;
		const unsigned p0 = _unsigned(p);
		Polynomial<ResidueField> root;

		for (unsigned i = 0; i < poly.size(); i += p0)
			root.push_back(poly[i]);

		const auto f = Factor(root);

		for (auto iter = f.begin(); iter != f.end(); ++iter)
			factor.push_back(std::make_pair(iter->first, iter->second * p0));

		std::sort(factor.begin(), factor.end());

		return factor;
	}

	std::vector<std::pair<Polynomial<ResidueField>, unsigned>> factor;
	Polynomial<ResidueField> f(squarefreepoly);
	Polynomial<ResidueField> h;
	h.push_back(ResidueField(0, p, true));
	h.push_back(ResidueField(1, p, true));
	const Polynomial<ResidueField> h0(h);

	for (unsigned deg = 1; Degree(f); ++deg)
	{
		h = std::Exp(PolynomialResidue<ResidueField>(h, squarefreepoly), p).Value();
		Polynomial<ResidueField> g = GCD(h - h0, f);
		f /= g;

		if (!Degree(g))
			continue;

		const auto fac = EqualDegreeFactor(g, deg);

		for (auto iter = fac.begin(); iter != fac.end(); ++iter)
			factor.push_back(std::make_pair(*iter, 0));
	}

	Polynomial<ResidueField> lcpoly(poly);

	for (auto iter = factor.begin(); iter != factor.end(); ++iter)
	{
		unsigned k = 0;

		do
		{
			++k;
			lcpoly /= iter->first;
		}
		while (lcpoly % iter->first == Polynomial<ResidueField>(0));

		*iter = std::make_pair(iter->first, k);
	}

	if (Degree(lcpoly))
	{
		const unsigned p0 = _unsigned(p);
		Polynomial<ResidueField> root;

		for (unsigned i = 0; i < lcpoly.size(); i += p0)
			root.push_back(lcpoly[i]);

		const auto f = Factor(root);

		for (auto iter = f.begin(); iter != f.end(); ++iter)
			factor.push_back(std::make_pair(iter->first, iter->second * p0));

		std::sort(factor.begin(), factor.end());

		return factor;
	}
	else
		return factor;
}

std::vector<std::pair<Polynomial<Integer>, unsigned>> Factor(const Polynomial<Integer> &poly)
{
	throw "TODO";
}

std::string NumberTheory::PrintFactor(const Integer &n)
{
	std::ostringstream str;
	const auto p = Factor(n);
	auto iter = p.begin();
	str << iter->first;

	if (iter->second > 1)
		str << "^" << iter->second;

	while (++iter != p.end())
	{
		str << " * " << iter->first;

		if (iter->second > 1)
			str << "^" << iter->second;
	}

	return str.str();
}

unsigned NumberTheory::DivisorCount(const Integer &n)
{
	if (n == 0)
		return 0;
	else if (n < 0)
		return DivisorCount(-n);
	else if (n == 1)
		return 1;

	unsigned count = 1;
	const auto p = Factor(n);

	for (auto iter = p.begin(); iter != p.end(); ++iter)
		count *= (iter->second + 1);

	return count;
}

Integer NumberTheory::DivisorSum(const Integer &n)
{
	if (n == 0 || n == 1)
		return n;

	Integer sum(1);
	const auto p = Factor(n);

	for (auto iter = p.begin(); iter != p.end(); ++iter)
		sum *= (std::Exp(iter->first, iter->second + 1) - 1) / (iter->first - 1);

	return sum;
}

Integer NumberTheory::EulerPhi(const Integer &n)
{
	const auto factor = Factor(n);
	Integer m(n);

	for (auto iter = factor.begin(); iter != factor.end(); ++iter)
		m -= m / iter->first;

	return m;
}

int NumberTheory::Mobius(const Integer &n)
{
	const auto factor = Factor(n);
	const unsigned size = factor.size();

	for (unsigned i = 0; i < size; ++i)
		if (factor[i].second > 1)
			return 0;

	return std::IsEven(factor.size()) ? 1 : -1;
}
